#!/usr/bin/env python
# -*- coding: utf-8 -*-

""" By Martin Senande-Rivera
    For Towards and atmosphere more favourable to firestorm development in Europe """

import os
import glob, sys
import numpy as np
import xarray as xr
import pandas as pd

# Fine Fuel Moisture Code calculation
def FFMCcalc(prcp,temp,rhum,wind,ffmc0):
	mo = (147.2*(101.0-ffmc0))/(59.5+ffmc0)
	rhum = xr.where(rhum>100.0, 100.0, rhum)
	rhum = xr.where(rhum<0.0, 0.0, rhum)
	wind = xr.where(wind<0.0, 0.0, wind)

	rf = prcp-0.5
	mo = xr.where((prcp>0.5) & (mo>150.), mo+42.5*rf*np.exp(-100.0/(251.0-mo))*(1.0-np.exp(-6.93/rf))+(.0015*(mo-150.0)**2)*np.sqrt(rf), mo)
	mo = xr.where((prcp>0.5) & (mo<=150.) & (mo<=250.), mo+42.5*rf*np.exp(-100.0/(251.0-mo))*(1.0-np.exp(-6.93/rf)), mo)
	mo = xr.where((prcp>0.5) & (mo>250.), 250., mo)

	ed = .942*(rhum**.679)+(11.0*np.exp((rhum-100.0)/10.0))+0.18*(21.1-temp)*(1.0-1.0/np.exp(.115*rhum))
	kl =.424*(1.0-(rhum/100.0)**1.7)+(.0694*np.sqrt(wind))*(1.0-(rhum/100.0)**8) 
	kd = kl*(.581*np.exp(.0365*temp)) 

	ew = .618*(rhum**.753)+(10.0*np.exp((rhum-100.0)/10.0))+0.18*(21.1-temp)*(1.0-1.0/np.exp(.115*rhum))
	kl = .424*(1.0-((100.0-rhum)/100.0)**1.7)+(.0694*np.sqrt(wind))*(1.0-((100.0-rhum)/100.0)**8) 
	kw = kl*(.581*np.exp(.0365*temp))

	m = xr.full_like(temp,np.nan)
	m = xr.where((mo<ed) & (mo<ew), ew-(ew-mo)/10.0**kw, m)
	m = xr.where((mo<ed) & (mo>=ew), mo, m)
	m = xr.where(mo==ed, mo, m)
	m = xr.where(mo>ed, ed+(mo-ed)/10.**kd, m)

	ffmc = (59.5*(250.0-m))/(147.2+m)
	ffmc = xr.where(ffmc>101.0, 101.0, ffmc)
	ffmc = xr.where(ffmc<=0.0, 0.0, ffmc)
	return ffmc

# Duff Moisture Code calculation
def DMCcalc(prcp,temp,rhum,dmc0,mth):
	el = [6.5,7.5,9.0,12.8,13.9,13.9,12.4,10.9,9.4,8.0,7.0,6.0]
	rhum = xr.where(rhum>100.0, 100.0, rhum)
	rhum = xr.where(rhum<0.0, 0.0, rhum)

	ra = prcp
	rw = 0.92*ra-1.27	
	wmi = 20.0+280.0/np.exp(0.023*dmc0)
	b = xr.full_like(temp,np.nan)
	b = xr.where((prcp>1.5) & (dmc0<=33.), 100.0/(0.5+0.3*dmc0), b)
	b = xr.where((prcp>1.5) & (dmc0<=65.) & (dmc0>33.), 14.0-1.3*np.log(dmc0), b)
	b = xr.where((prcp>1.5) & (dmc0>65.), 6.2*np.log(dmc0)-17.2, b)
	wmr = wmi+(1000*rw)/(48.77+b*rw)

	pr = xr.full_like(temp,np.nan)
	pr = xr.where(prcp>1.5, 43.43*(5.6348-np.log(wmr-20.0)), pr)
	pr = xr.where(prcp<=1.5, dmc0, pr)
	pr = xr.where(pr<0.0, 0.0, pr)

	rk = xr.full_like(temp,np.nan)
	rk = xr.where(temp>=-1.1, 1.894*(temp+1.1)*(100.0-rhum)*(el[mth-1]*0.0001), rk) 
	rk = xr.where(temp<-1.1, 0., rk) 
	
	dmc = pr+rk
	dmc = xr.where(dmc<=0.0, 0.0, dmc)
	return dmc	

# Drought Code calculation
def DCcalc(prcp,temp,dc0,mth):
	fl = [-1.6, -1.6, -1.6, 0.9, 3.8, 5.8, 6.4, 5.0, 2.4, 0.4, -1.6, -1.6]	

	rd = 0.83*prcp-1.27
	q0 = 800.0*np.exp(-dc0/400.0)
	qr = q0+3.937*rd
	dr = 400.0*np.log(800.0/qr)
	dr = xr.where(dr<=0.0, 0.0, dr)

	v = xr.full_like(temp,np.nan)
	v = xr.where(temp>=-2.8, 0.36*(temp+2.8)+fl[mth-1], v)
	v = xr.where(temp<-2.8, fl[mth-1], v)
	v = xr.where(v<0.0, 0.0, v)
	dc = xr.full_like(temp,np.nan)
	dc = xr.where(prcp>2.8, dr+0.5*v, dc)
	dc = xr.where(prcp<=2.8, dc0+0.5*v, dc)
	return dc

# Initial Spread Index calculation
def ISIcalc(wind,ffmc):
	mo = 147.2*(101.0-ffmc)/(59.5+ffmc)
	ff = 19.115*np.exp(mo*(-0.1386))*(1.0+(mo**5.31)/49300000.0)
	isi = ff*np.exp(0.05039*wind)
	return isi

# Buildup Index calculation
def BUIcalc(dmc,dc):
	bui = xr.full_like(dmc,np.nan)
	bui = xr.where(dmc<=0.4*dc, (0.8*dc*dmc)/(dmc+0.4*dc), bui)
	bui = xr.where(dmc>0.4*dc, dmc-(1.0-0.8*dc/(dmc+0.4*dc))*(0.92+(0.0114*dmc)**1.7), bui)
	bui = xr.where((dmc==0.0) & (dc==0.0), 0.0, bui)
	bui = xr.where(bui<0.0, 0.0, bui)
	return bui

# Fire Weather Index calculation
def FWIcalc(isi,bui):
	bb = xr.full_like(bui,np.nan)
	bb = xr.where(bui<=80.0, 0.1*isi*(0.626*bui**0.809+2.0), bb)
	bb = xr.where(bui>80.0, 0.1*isi*(1000.0/(25.+108.64/np.exp(0.023*bui))), bb)
	
	fwi = xr.full_like(bui,np.nan)
	fwi = xr.where(bb<=1.0, bb, fwi)
	fwi = xr.where(bb>1.0, np.exp(2.72*(0.434*np.log(bb))**0.647), fwi)
	fwi = xr.where(fwi<=1.0, 1.0, fwi)
	return fwi

GCM=sys.argv[1]        # Enter GCM
RCM=sys.argv[2]        # Enter RCM
rcp=sys.argv[3]        # Enter RCP scenario
year=int(sys.argv[4])  # Enter year

path_rh='./EURO-CORDEX/'+RCM+'/'+GCM+'/hurs/'     # Relative humidity in %
path_ta='./EURO-CORDEX/'+RCM+'/'+GCM+'/tas/'      # 2m air temperature in Kelvin
path_ws='./EURO-CORDEX/'+RCM+'/'+GCM+'/sfcWind/'  # Surface wind speed in m s-1
path_tp='./EURO-CORDEX/'+RCM+'/'+GCM+'/pr/'       # Precipitation flux in kg m-2 s-1
path_lm='./EURO-CORDEX/'+RCM+'/'+GCM+'/lm/'       # Landmask

path_out_ffmc='./FWI/'+RCM+'/'+GCM+'/FFMC/'
path_out_dmc='./FWI/'+RCM+'/'+GCM+'/DMC/'
path_out_dc='./FWI/'+RCM+'/'+GCM+'/DC/'
path_out_fwi='./FWI/'+RCM+'/'+GCM+'/FWI/'

if (GCM=='EC-EARTH'):
	GCM2='ICHEC-EC-EARTH'
	dia2='31'
elif (GCM=='HadGEM2-ES'):
	GCM2='MOHC-HadGEM2-ES'
	dia2='30'
elif (GCM=='MPI-ESM-LR'):
	GCM2='MPI-M-MPI-ESM-LR'
	dia2='31'
elif (GCM=='NorESM1-M'):
	GCM2='NCC-NorESM1-M'
	dia2='31'

if year<2006:
	experiment='historical'
else:
	experiment=rcp

## Landmask
lm = xr.open_dataset(path_lm+'landmask.nc')['sftlf']
## Relative humidity
rh = xr.open_dataset(path_rh+'hurs_EUR-11_'+GCM2+'_'+experiment+'_12hr_'+str(year)+'.nc',decode_times=False)['hurs']
## 2m air temperature
ta = xr.open_dataset(path_ta+'tas_EUR-11_'+GCM2+'_'+experiment+'_12hr_'+str(year)+'.nc',decode_times=False)['tas']-273.15
## Surface wind speed
ws = xr.open_dataset(path_ws+'sfcWind_EUR-11_'+GCM2+'_'+experiment+'_12hr_'+str(year)+'.nc',decode_times=False)['sfcWind']*3.6
## Total precipitation
tp = xr.open_dataset(path_tp+'pr_EUR-11_'+GCM2+'_'+experiment+'_'+str(year)+'.nc',decode_times=False)['pr']*86400.*0.997

# Creating time vector
if (GCM=='EC-EARTH'):
	dr = pd.date_range(start=str(year)+'-01-01', end=str(year)+'-12-31', freq='1D')
	dates = dr
elif (GCM=='HadGEM2-ES'):
	dr = pd.date_range(start=str(year)+'-01-01', end=str(year)+'-12-31', freq='1D')
	dates = dr[((dr.day != 29) | (dr.month != 2)) & ((dr.day != 31) | (dr.month <4))]
elif (GCM=='MPI-ESM-LR'):
	dr = pd.date_range(start=str(year)+'-01-01', end=str(year)+'-12-31', freq='1D')
	dates = dr
elif (GCM=='NorESM1-M'):
	dr = pd.date_range(start=str(year)+'-01-01', end=str(year)+'-12-31', freq='1D')
	dates = dr[(dr.day != 29) | (dr.month != 2)]
rh['time'] = dates
ta['time'] = dates
ws['time'] = dates
tp['time'] = dates

# Reading moisture codes of day before
if year==1970:
	ffmc0 = xr.full_like(ta.isel(time=0), 85.0)
	#ffmc0 = xr.where(lm==0.,np.nan,ffmc0)
	dmc0 = xr.full_like(ta.isel(time=0), 6.0)
	#dmc0 = xr.where(lm==0.,np.nan,dmc0)
	dc0 = xr.full_like(ta.isel(time=0), 15.0)
	#dc0 = xr.where(lm==0.,np.nan,dc0)
elif year==2006:
	ffmc0 = xr.open_dataset(path_out_ffmc+'ffmc_historical_'+str(year-1)+'.nc',decode_times=False)['ffmc'].copy()
	dmc0 = xr.open_dataset(path_out_dmc+'dmc_historical_'+str(year-1)+'.nc',decode_times=False)['dmc'].copy()
	dc0 = xr.open_dataset(path_out_dc+'dc_historical_'+str(year-1)+'.nc',decode_times=False)['dc'].copy()
else:
	ffmc0 = xr.open_dataset(path_out_ffmc+'ffmc_'+experiment+'_'+str(year-1)+'.nc',decode_times=False)['ffmc'].copy()
	dmc0 = xr.open_dataset(path_out_dmc+'dmc_'+experiment+'_'+str(year-1)+'.nc',decode_times=False)['dmc'].copy()
	dc0 = xr.open_dataset(path_out_dc+'dc_'+experiment+'_'+str(year-1)+'.nc',decode_times=False)['dc'].copy()

# Loop for computing daily FWI
fwi = xr.full_like(ta,np.nan)
for t in range(ta.time.size):
	print(dates[t])

	mth = dates[t].month
	temp = ta.isel(time=t).copy()
	rhum = rh.isel(time=t).copy()
	wind = ws.isel(time=t).copy()
	prcp = tp.isel(time=t).copy()
	
	ffmc = FFMCcalc(prcp,temp,rhum,wind,ffmc0)
	dmc = DMCcalc(prcp,temp,rhum,dmc0,mth)
	dc = DCcalc(prcp,temp,dc0,mth)
	isi = ISIcalc(wind,ffmc)
	bui = BUIcalc(dmc,dc)
	fwi0 = FWIcalc(isi,bui)
	fwi0 = xr.where(lm==0.,np.nan,fwi0)
	fwi.values[t,:] = fwi0.copy()
	ffmc0 = ffmc.copy()
	dmc0 = dmc.copy()
	dc0 = dc.copy()


print('Year: {} \n Loop end. Saving files...'.format(str(year)))

ffmc0.name = 'ffmc'
ffmc0.attrs['standard_name'] = 'ffmc'
ffmc0.attrs['long_name'] = 'Fine Fuel Moisture Code'
ffmc0.attrs['units'] = '1'

dmc0.name = 'dmc'
dmc0.attrs['standard_name'] = 'dmc'
dmc0.attrs['long_name'] = 'Duff Moisture Code'
dmc0.attrs['units'] = '1'

dc0.name = 'dc'
dc0.attrs['standard_name'] = 'dc'
dc0.attrs['long_name'] = 'Drought Code'
dc0.attrs['units'] = '1'

fwi.name = 'fwi'
fwi.attrs['standard_name'] = 'fwi'
fwi.attrs['long_name'] = 'Fire Weather Index'
fwi.attrs['units'] = '1'

# Saving files
ffmc0.to_netcdf(path_out_ffmc+'ffmc_'+experiment+'_'+str(year)+'.nc')
dmc0.to_netcdf(path_out_dmc+'dmc_'+experiment+'_'+str(year)+'.nc')
dc0.to_netcdf(path_out_dc+'dc_'+experiment+'_'+str(year)+'.nc')
fwi.to_netcdf(path_out_fwi+'fwi_'+experiment+'_'+str(year)+'.nc')
